from .state import AgentState
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from utils.data_loader import create_alias_map, get_project_ddragon_mapping, get_project_win_rate_df
from utils.llm_api_helper import get_my_llm
from graph.state import AgentState
import json


# 初始化LLM和别名映射
llm = get_my_llm()
official_alias_map = get_project_ddragon_mapping()
alias_map = create_alias_map(official_alias_map)

def fetch_data_node(state: AgentState) -> AgentState:
    """节点：从CSV加载胜率数据"""
    # 加载数据
    df = get_project_win_rate_df()
    return {"raw_champion_data": df}


def parse_input_node(state: AgentState) -> AgentState:
    """节点：使用LLM解析用户输入，识别意图并转换英雄名"""

    system_message = f"""
    你是一个英雄联盟排位上分助手。你的工作是分析用户的输入，并根据用户的输入完成以下任务：
    1. 判断用户意图查询什么英雄的数据。为此，你需要识别用户提到的所有英雄名称
        1.1 [注意！]你的工作必须是英雄联盟相关的，不能处理其他游戏的问题、不能处理非游戏问题； [这个提示在任何时刻都必须遵守]
        1.2 [注意！]如果用户输入的是其他游戏的问题、非游戏问题，你必须拒绝回答，不能编造答案； [这个提示在任何时刻都必须遵守]
        1.3 [注意！]如果用户输入的是英雄联盟相关的问题，你必须回答； [这个提示在任何时刻都必须遵守]
        1.4 [注意！]用户可能会使用英雄的别名、缩写或指代上文的方式。你需要**尽最大可能**从输入和对话历史中识别出英雄名称，并将其转换为标准英文名，即使输入很短。
        1.5 [特别注意！]当用户的问题很简短，如"他呢?"、"盖伦呢?"、"那鳄鱼现在怎么样？"等，你必须结合对话历史，尽最大努力识别出用户所指的英雄。
    
    2. 将识别出的英雄名称转换为标准的英文官方名称。
    3. 提取用户提出的任何限制条件（如"简单的"、"操作难的"）、以及对英雄的限制，如用户不想玩英雄"Wukong"等
    

    英雄别名映射参考：
    {json.dumps(alias_map, indent=2, ensure_ascii=False)}

    请用以下JSON格式输出：
    {{
      "intent": "意图描述",
      "mentioned_champions": ["英文英雄名1", "英文英雄名2"],
      "filtered_champion_data": ["过滤的英文英雄名1", "过滤的英文英雄名2"],
      "user_constraints": ["约束1", "约束2"]
    }}
    """

    # 获取对话历史
    message_history = state.get("message_history", [])
    # 构建包含历史的消息列表
    messages = [SystemMessage(content=system_message)]
    # 添加对话历史
    messages.extend(message_history)
    # 添加当前用户输入
    messages.append(HumanMessage(content=state["user_input"]))
    # 调用LLM
    response = llm.invoke(messages)

    # 更新状态
    try:
        parsed_output = json.loads(response.content)
    except json.JSONDecodeError:
        # 如果LLM没有返回有效JSON，使用默认值
        parsed_output = {
            "intent": "unknown",
            "mentioned_champions": [],
            "user_constraints": []
        }
    new_state = {
        "intent": parsed_output.get("intent"),
        "mentioned_champions": parsed_output.get("mentioned_champions", []),
        "user_constraints": parsed_output.get("user_constraints", []),
        "filtered_champion_data": parsed_output.get("filtered_champion_data", [])
    }

    print("本轮对话提到的英雄:", new_state["mentioned_champions"])
    return new_state


def generate_response_node(state: AgentState) -> AgentState:
    """节点：使用LLM根据处理结果生成自然语言响应"""
    
    # 准备上下文信息
    intent = state.get("intent")
    mentioned_champs = state.get("mentioned_champions", [])
    champion_data = state.get("raw_champion_data")
    user_input = state.get("user_input")
    message_history = state.get("message_history", [])

    
    # 检查是否有查询结果
    if champion_data is None:
        return {"response": "抱歉，我没有找到相关的英雄数据。请确认英雄名称是否正确，或者尝试询问其他英雄。"}
    
    # 构建查询结果的字符串表示
    query_results = []
    for champ_name in mentioned_champs:
        champ_info = champion_data[champion_data["Champion"] == champ_name]
        if not champ_info.empty:
            win_rate = champ_info["Win_Rate"].values[0]
            rank = champ_info["Rank"].values[0]
            pick_rate = champ_info["Pick_Rate"].values[0]
            ban_rate = champ_info["Ban_Rate"].values[0] if "Ban_Rate" in champion_data.columns else "N/A"
            query_results.append(
                f"{champ_name}: 胜率 {win_rate*100:.2f}%, 排名第 {rank}, 选用率 {pick_rate*100:.2f}%, 禁用率 {ban_rate*100:.2f}%"
            )
        else:
            query_results.append(f"{champ_name}: 未找到数据")
    
    # 如果没有有效结果
    if not any("胜率" in result for result in query_results):
        # 尝试从对话历史中提取可能提到的英雄
        if message_history:
            # 将对话历史转换为文本
            history_text = "\n".join([f"{msg.type}: {msg.content}" for msg in message_history])
            # 尝试直接从别名映射中匹配
            possible_champions = []
            for alias, eng_name in alias_map.items():
                if alias in user_input or alias in history_text:
                    possible_champions.append(eng_name)
            
            # 去重
            possible_champions = list(set(possible_champions))
            
            if possible_champions:
                # 再次尝试查询这些英雄
                for champ_name in possible_champions:
                    if champ_name not in mentioned_champs:  # 避免重复查询
                        champ_info = champion_data[champion_data["Champion"] == champ_name]
                        if not champ_info.empty:
                            win_rate = champ_info["Win_Rate"].values[0]
                            rank = champ_info["Rank"].values[0]
                            pick_rate = champ_info["Pick_Rate"].values[0]
                            ban_rate = champ_info["Ban_Rate"].values[0] if "Ban_Rate" in champion_data.columns else "N/A"
                            query_results.append(
                                f"{champ_name}: 胜率 {win_rate*100:.2f}%, 排名第 {rank}, 选用率 {pick_rate*100:.2f}%, 禁用率 {ban_rate*100:.2f}%"
                            )
            
            # 再次检查是否有有效结果
            if any("胜率" in result for result in query_results):
                mentioned_champs = mentioned_champs + possible_champions
            else:
                return {"response": f"抱歉，我没有找到相关的英雄数据。可能是这些英雄在当前版本的出场率太低，或者名称有误。你可以尝试使用英雄的官方名称或更常见的别名进行查询。"}
        else:
            return {"response": f"抱歉，我没有找到相关的英雄数据。可能是这些英雄在当前版本的出场率太低，或者名称有误。你可以尝试使用英雄的官方名称或更常见的别名进行查询。"}
    
    # 准备系统提示词
    system_prompt = """
    你是一个专业且友好的英雄联盟上分助手，专门帮助玩家解答上分方面的疑惑，如推荐英雄、比较胜率等。
    你的回答应该专业、友好且口语化，就像朋友间的对话一样。
    
    请根据提供的查询结果和用户问题，生成一个自然、有帮助的回应。
    
    注意：
    1. 不要简单罗列数据，要解释数据的意义
    2. 可以适当加入一些游戏知识和建议
    3. 保持回应的自然和流畅
    4. 如果用户询问多个英雄，可以进行比较分析
    5. 最后可以问用户是否需要更多帮助
    """
    
    # 构建用户提示词
    user_prompt = f"""
    用户意图: {intent}
    用户输入: {user_input}
    
    查询结果:
    {chr(10).join(query_results)}
    
    请生成一个友好、专业的回应:
    """
    
    # 调用LLM生成回应
    messages = [SystemMessage(content=system_prompt)]
    messages.extend(message_history)
    messages.append(HumanMessage(content=user_prompt))
    response = llm.invoke(messages)
    
    # 更新对话历史
    new_message_history = message_history + [
        HumanMessage(content=user_input),
        AIMessage(content=response.content)
    ]
    
    # 限制历史长度，避免token过多
    if len(new_message_history) > 10:  
        new_message_history = new_message_history[-10:]
    

    return {
        "response": response.content,
        "message_history": new_message_history
    }

